package com.bupt.ilink.controller;

import com.bupt.ilink.entity.VolcText2ImgRequest;
import com.bupt.ilink.entity.VolcImg2ImgRequest;
import com.bupt.ilink.utils.R;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.JSONArray;
import com.volcengine.service.visual.IVisualService;
import com.volcengine.service.visual.impl.VisualServiceImpl;
import com.volcengine.service.visual.model.response.VisualHighAesSmartDrawingResponse;

import java.io.File;
import java.io.FileOutputStream;
import java.io.InputStream;
import java.net.HttpURLConnection;
import java.net.URL;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;


@RestController
@RequestMapping("/volcengine")
public class VolcEngineController {
    /**
     *
     * @Author Ziyan Ma
     * @Description 火山引擎-豆包-图生文大模型接口
     * @Param request 见VolcText2ImgRequest参数
     * @return R响应
     **/
    @PostMapping("/text2img")
    public R text2Img(@RequestBody VolcText2ImgRequest request) {
        System.out.println("START START START");
        IVisualService visualService = VisualServiceImpl.getInstance();
        visualService.setAccessKey(request.getAk());
        visualService.setSecretKey(request.getSk());
        JSONObject req = new JSONObject();
        req.put("req_key", "high_aes_general_v20");
        req.put("model_version", "general_v2.0");
        req.put("prompt", request.getPrompt());
        req.put("return_url", true);
        try {
            VisualHighAesSmartDrawingResponse response = visualService.visualHighAesSmartDrawing(req);
            JSONObject jsonObject = (JSONObject) JSON.toJSON(response);
            return R.ok("", jsonObject.getJSONObject("data").getJSONArray("image_urls").getString(0));
        } catch (Exception e) {
            return R.failure(e.getMessage());
        }
    }

    /**
     *
     * @Author Ziyan Ma
     * @Description 火山引擎-豆包-图生图大模型接口
     * @Param request 见VolcImg2ImgRequest参数
     * @return R响应
     **/
    @PostMapping("/img2img")
    public R img2Img(@RequestBody VolcImg2ImgRequest request) throws Exception {
        IVisualService visualService = VisualServiceImpl.getInstance();
        visualService.setAccessKey(request.getAk());
        visualService.setSecretKey(request.getSk());
        JSONObject req = new JSONObject();
        System.out.println(request.getImages());
        List<String> images = convertImagesToBase64(request.getImages());
        System.out.println('>');
        System.out.println(images );
        // 必选参数添加
        req.put("binary_data_base64", images);
        req.put("req_key", "i2i_xl_sft");
        req.put("return_url", true);

        JSONObject controlnetArg = new JSONObject();
        controlnetArg.put("type", "canny");
        controlnetArg.put("strength", 0.4);
        controlnetArg.put("binary_data_index", 0);

        JSONArray controlnetArgs = new JSONArray();
        controlnetArgs.add(controlnetArg);
        req.put("controlnet_args", controlnetArgs);

        JSONObject styleReferenceArgs = new JSONObject();
        styleReferenceArgs.put( "binary_data_index", 0);
        req.put("style_reference_args",  styleReferenceArgs);

        try{
            Object response = visualService.img2ImgXLSft(req);
//            System.out.println(response);
            JSONObject resContent = (JSONObject) JSON.toJSON(response);
            return R.ok("", resContent.getJSONObject("data").getJSONArray("image_urls").getString(0));
        } catch (Exception e) {
            return R.failure(e.getMessage());
        }
    }


    /**
     *
     * @Author Ziyan Ma
     * @Description 图像url转化为base64格式
     * @Param imageUrls 列表
     * @return 转化后的base64 列表
     **/
    public static List<String> convertImagesToBase64(List<String> imageUrls) throws Exception {
        List<String> base64ImageList = new ArrayList<>();

        for (String imageUrl : imageUrls) {
            File tempFile = null;
            try {
                // Download the image
                URL url = new URL(imageUrl);
                HttpURLConnection connection = (HttpURLConnection) url.openConnection();
                connection.setRequestMethod("GET");
                InputStream inputStream = connection.getInputStream();

                // Save image to a temporary file
                tempFile = File.createTempFile("image", ".tmp");
                try (FileOutputStream fileOutputStream = new FileOutputStream(tempFile)) {
                    byte[] buffer = new byte[1024];
                    int bytesRead;
                    while ((bytesRead = inputStream.read(buffer)) != -1) {
                        fileOutputStream.write(buffer, 0, bytesRead);
                    }
                }

                // Convert image file to byte array
                byte[] imageBytes = java.nio.file.Files.readAllBytes(tempFile.toPath());

                // Encode byte array to Base64
                String base64Image = Base64.getEncoder().encodeToString(imageBytes);
                base64ImageList.add(base64Image);

                // Close streams
                inputStream.close();
            } catch (Exception e) {
                e.printStackTrace(); // Handle the exception appropriately
            } finally {
                // Delete the temporary file
                if (tempFile != null && tempFile.exists()) {
                    tempFile.delete();
                }
            }
        }
        return base64ImageList;
    }
}
